import math

import numpy as np

import layer_utils


def nms(args, classes, offsets, anchors):
    objects = np.argmax(classes, axis=1)
    nonbg = np.nonzero(objects)[0]

    if args.no_img:
        nonbg = nonbg[:10]

    indexes = []
    while True:
        scores = np.zeros((classes.shape[0],))
        scores[nonbg] = np.amax(classes[nonbg], axis=1)

        score_idx = np.argmax(scores, axis=0)
        score_max = scores[score_idx]

        # 去掉
        nonbg = nonbg[nonbg != score_idx]

        if score_max < args.class_threshold:
            break

        indexes.append(score_idx)
        score_anc = anchors[score_idx]
        score_off = offsets[score_idx][0:4]
        score_box = score_anc + score_off

        score_box = np.expand_dims(score_box, axis=0)

        nonbg_copy = nonbg.copy()
        for idx in nonbg_copy:
            anchor = anchors[idx]
            offset = offsets[idx][:4]
            box = anchor + offset
            box = np.expand_dims(box, axis=0)

            iou = layer_utils.iou(box, score_box)[0][0]

            if args.soft_nms:
                iou = - 2 * iou * iou
                classes[idx] *= math.exp(iou)

            elif iou >= args.iou_threshold:
                nonbg = nonbg[nonbg != idx]

        if nonbg.size == 0:
            break
    scores = np.zeros((classes.shape[0]))
    scores[indexes] = np.amax(classes[indexes], axis=1)
    return objects, indexes, scores


def get_nms(args, image, classes, offsets, feature_shapes):
    anchors = []
    for index, feature_shape in enumerate(feature_shapes):
        anchor = layer_utils.anchor_boxes(feature_shape,
                                          image.shape,
                                          index=index)
        anchor = np.reshape(anchor, [-1, 4])
        if index == 0:
            anchors = anchor
        else:
            anchors = np.concatenate((anchors, anchor), axis=0)

    anchors_centroid = layer_utils.minmax2centroid(anchors)
    offsets[:, 0:2] *= 0.1
    offsets[:, 0:2] *= anchors_centroid[:, 2:4]
    offsets[:, 0:2] += anchors_centroid[:, 0:2]
    offsets[:, 2:4] *= 0.2
    offsets[:, 2:4] = np.exp(offsets[:, 2:4])
    offsets[:, 2:4] *= anchors_centroid[:, 2:4]
    offsets = layer_utils.centroid2minmax(offsets)
    # convert fr cx,cy,w,h to real offsets
    offsets[:, 0:4] = offsets[:, 0:4] - anchors

    objects, indexes, scores = nms(args, classes, offsets, anchors)
    boxes = []
    for idx in indexes:
        # batch, row, col, box
        anchor = anchors[idx]
        offset = offsets[idx]

        anchor += offset[0:4]

        boxes.append(anchor)
    return objects, indexes, scores, boxes